package com.sys.ai.workflow;

import com.agentsflex.core.chain.Chain;
import com.agentsflex.core.document.Document;
import com.sys.ai.domain.AiKnowledge;
import com.sys.ai.service.AiService;
import com.sys.ai.service.IAiKnowledgeService;
import com.sys.ai.service.IDbOptionService;
import com.sys.ai.vo.EmbeddingsVO;
import com.sys.ai.vo.KnowledgeFragmentationVO;
import dev.tinyflow.core.knowledge.Knowledge;
import dev.tinyflow.core.node.KnowledgeNode;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;

/**
 * @author LGY
 * Create by 2025/5/26 8:58
 */
@Service
public class CustomerKnowledge implements Knowledge {

    @Autowired
    private IDbOptionService embeddingDBService;
    @Autowired
    private AiService aiService;
    @Autowired
    private IAiKnowledgeService aiKnowledgeService;

    @Override
    public List<Document> search(String keyword, int limit, KnowledgeNode knowledgeNode, Chain chain) {
        String knowledgeId = knowledgeNode.getKnowledgeId().toString();
        AiKnowledge aiKnowledge = aiKnowledgeService.getById(knowledgeId);
        // 将问题转成向量
        EmbeddingsVO embeddingsVO = aiService.generateQAToEmbeddings(keyword, aiKnowledge.getRankModelId());
        List<KnowledgeFragmentationVO> knowledgeQAVOList = embeddingDBService.query(embeddingsVO.getData().get(0).getEmbedding(), List.of(knowledgeId));
        return knowledgeQAVOList.stream().map(item -> {
            Document document = new Document();
            String content = item.getFragmentation();
            document.setId(aiKnowledge.getVectorDb());
            document.setContent(content);
            document.setTitle(aiKnowledge.getName());
            return document;
        }).toList();
    }
}
